
/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 *
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

#ifndef ATVC_BROADCAST_OP_TEMPLATE_H
#define ATVC_BROADCAST_OP_TEMPLATE_H
#include "common/const_def.h"
#include "common/ops_utils_device.h"
#include "broadcast/broadcast_utils/broadcast_buf_pool.h"
#include "broadcast/broadcast_compute.h"
namespace ATVC {
struct BroadcastDataView {
    uint32_t dimASize;
    uint32_t dimBSize;
    uint32_t inShape[ATVC::MAX_DIM];
    uint32_t outShape[ATVC::MAX_DIM];
    uint32_t copyInSize;        // 单核拷入数据量
    uint32_t A11;               // 实际参与计算的A11
    uint32_t A12;               // 实际参与计算的A12
    uint32_t B1;                // 实际参与计算的B1
    uint32_t dimAOffset;        // 输入输出数据在A维度的偏移量
    uint32_t dimBOffset;        // 输入输出数据在B维度的偏移量
    uint32_t copyOutBaseOffset; // 核间数据拷出基址
};

namespace Kernel {
template <class BroadcastCompute, const auto& SelectBroadcastPolicy, typename ParamPtr = __gm__ ATVC::BroadcastParam*>
class BroadcastOpTemplate {
public:
    using DataType = typename BroadcastCompute::DataType;
    __aicore__ inline BroadcastOpTemplate() {}
    /*
    BroadcastOpTemplate对外运行接口，主要完成资源初始化、数据搬入、计算调度、数据搬出操作
    @param src: 输入数据的gm指针
    @param dst: 输出数据的gm指针
    @broadcastParam: broadcast的动态参数，包含tiling data, workspace等
    */
    template<typename T1>
    __aicore__ inline void Run(GM_ADDR src, GM_ADDR dst, T1 broadcastParam)
    {
        this->Init(src, dst, broadcastParam);
        this->Process();
    }

private:
    template<typename T1>
    __aicore__ inline void Init(GM_ADDR src, GM_ADDR dst, T1 broadcastParam)
    {
        param_ = reinterpret_cast<ParamPtr>(broadcastParam);
        uint32_t srcDataSize = this->param_->tilingData.basicBlock;
        uint32_t dstDataSize = this->param_->tilingData.basicBlock;
        srcGlobal_.SetGlobalBuffer(reinterpret_cast<__gm__ DataType*>(src), srcDataSize);
        dstGlobal_.SetGlobalBuffer(reinterpret_cast<__gm__ DataType*>(dst), dstDataSize);
        bufPool_.template Init<DataType>(GetTPipePtr(),
                                    ATVC::CONST2, // doublebuff需要的输入个数
                                    ATVC::CONST2, // 计算结果的个数，一般与inputNum保持一致
                                    this->param_->tilingData.A2 * this->param_->tilingData.A12 * DATA_SIZE, // 输入Tensor大小
                                    this->param_->tilingData.A2 * this->param_->tilingData.B2 * DATA_SIZE); // 输出Tensor大小
    }

    __aicore__ inline void CopyOutBatch(BroadcastDataView &view,
                                        uint32_t dimACount,
                                        AscendC::LocalTensor<DataType> &output)
    {
        uint32_t dimBCount = 0;          
        SyncDataQueue<AscendC::HardEvent::V_MTE3>(); 
        for (int i = 0; i < view.B1; i++) {
            uint32_t copyOutOffset;
            if (SelectBroadcastPolicy.patternID == AB_PATTERN::ABA) {
                copyOutOffset = dimBCount * view.dimASize + dimACount * this->param_->tilingData.A2;
            } else {
                copyOutOffset = dimACount * this->param_->tilingData.A2 * view.dimBSize + dimBCount;
            }
            CopyOut(output, copyOutOffset + view.copyOutBaseOffset, view);
            dimBCount += this->param_->tilingData.B2;
            AscendC::PipeBarrier<PIPE_MTE3>();
        }
    }

    __aicore__ inline void Process()
    {
        BroadcastDataView view;
        CalcView(view);
        uint32_t inputOffset;
        uint32_t dimACount = 0;
        AscendC::LocalTensor<DataType> input;
        for (int i = 0; i < view.A11; i++) {
            inputOffset = 0;
            bufPool_.AllocTensor<true>(input);
            uint32_t copyInOffset = i * view.A12 * this->param_->tilingData.A2;
            if (this->param_->tilingData.A0 != 1) {
                copyInOffset += view.dimAOffset;
            }
            if (copyInOffset >= view.dimASize) {
                return;
            }
            if (copyInOffset + view.copyInSize > view.dimASize) {
                // 剩下的数据不够一次完整计算， 根据实际数据重新计算
                view.copyInSize = view.dimASize - copyInOffset;
                view.A12 = OpsUtils::CeilDiv<uint32_t>(view.copyInSize, this->param_->tilingData.A2);
            }
            CopyIn(input, copyInOffset, view);
            bufPool_.SetVecSync(input);
            bufPool_.WaitVecSync(input);
            for (int j = 0; j < view.A12; j ++) {
                AscendC::LocalTensor<DataType> output;
                bufPool_.AllocTensor<false>(output);
                SyncDataQueue<AscendC::HardEvent::MTE2_V>();
                compute_.template Compute<SelectBroadcastPolicy.patternID>(input, inputOffset, output,
                    OpsUtils::CeilAlign<uint32_t>(this->param_->tilingData.A2, UB_ALIGN_COUNT),
                    OpsUtils::CeilAlign<uint32_t>(this->param_->tilingData.B2, UB_ALIGN_COUNT));
                bufPool_.SetCopyOutSync(output);
                bufPool_.WaitCopyOutSync(output);
                CopyOutBatch(view, dimACount, output);
                bufPool_.FreeTensor<false>(output);
                dimACount++;
                inputOffset += this->param_->tilingData.A2;
                SyncDataQueue<AscendC::HardEvent::MTE3_V>();
            }
            SyncDataQueue<AscendC::HardEvent::MTE3_MTE2>();
            bufPool_.FreeTensor<true>(input);            
        }
        bufPool_.ResetEvent();
    }

    __aicore__ inline uint32_t CalcCopyOutBaseOffset(BroadcastDataView &view)
    {
        uint32_t copyOutBaseOffset = 0;
        // 计算拷出偏移基址
        if (SelectBroadcastPolicy.patternID == AB_PATTERN::ABA) {
            if (this->param_->tilingData.A0 != 1) { // 核间A切分， 取部分A
                copyOutBaseOffset += view.dimAOffset;
            }
            if (this->param_->tilingData.B0 != 1) { // 核间B切分，取部分B
                copyOutBaseOffset += view.dimBOffset * view.dimASize;
            }
        } else {
            if (this->param_->tilingData.A0 != 1) { // 核间A切分， 取部分A
                copyOutBaseOffset += view.dimAOffset * view.dimBSize;
            }
            if (this->param_->tilingData.B0 != 1) { // 核间B切分，取部分B
                copyOutBaseOffset += view.dimBOffset;
            }
        }    
        return copyOutBaseOffset;    
    }

    __aicore__ inline void CalcView(BroadcastDataView &view)
    {
        if (SelectBroadcastPolicy.patternID == AB_PATTERN::ABA) {
            view.dimASize = this->param_->tilingData.dstShape[1];
            view.dimBSize = this->param_->tilingData.dstShape[0];
            view.inShape[0] = 1;
            view.inShape[1] = this->param_->tilingData.A2;
            view.outShape[0] = this->param_->tilingData.B2;
            view.outShape[1] = this->param_->tilingData.A2;
        } else {
            view.dimASize = this->param_->tilingData.dstShape[0];
            view.dimBSize = this->param_->tilingData.dstShape[1];
            view.inShape[0] = this->param_->tilingData.A2;
            view.inShape[1] = 1;
            view.outShape[0] = this->param_->tilingData.A2;
            view.outShape[1] = this->param_->tilingData.B2;
        }        
        view.A11 = this->param_->tilingData.A11;
        view.A12 = this->param_->tilingData.A12;
        view.B1 = this->param_->tilingData.B1;
        uint32_t blockId = AscendC::GetBlockIdx();
        uint32_t dimAIdx = blockId / this->param_->tilingData.B0;
        uint32_t dimBIdx = blockId % this->param_->tilingData.factorBTotalCnt;
        view.dimAOffset = dimAIdx * this->param_->tilingData.factorACntPerCore;
        view.dimBOffset = dimBIdx * this->param_->tilingData.factorBCntPerCore;
        // 计算一次计算的输入数据大小
        view.copyInSize = view.A12 * this->param_->tilingData.A2; // 一次拷贝A12份数据， for循环计算A12次
        if (view.dimAOffset + this->param_->tilingData.factorACntPerCore > view.dimASize) {
            // 剩下的A维度的数据不够每个核分到的A数目，重新计算实际的A维度切分
            uint32_t realShape = view.dimASize - view.dimAOffset;
            uint32_t A1 =  OpsUtils::CeilDiv<uint32_t>(realShape, this->param_->tilingData.A2);
            if (A1 < view.A12) {
                view.A11 = 1;
                view.A12 = A1;
            } else {
                view.A11 = OpsUtils::CeilDiv<uint32_t>(A1, view.A12);
            }
        }
        if (view.dimBOffset + this->param_->tilingData.factorBCntPerCore > view.dimBSize) {
            uint32_t realShape = view.dimBSize - view.dimBOffset;
            view.B1 =  OpsUtils::CeilDiv<uint32_t>(realShape, this->param_->tilingData.B2);
        }
        view.copyOutBaseOffset = CalcCopyOutBaseOffset(view);
    }

    __aicore__ inline void CopyIn(AscendC::LocalTensor<DataType> &input, uint32_t copyInOffset, BroadcastDataView &view)
    {
        AscendC::DataCopyPadExtParams<DataType> padParams{false, 0, 0, 0};
        AscendC::DataCopyExtParams copyInParams;
        copyInParams.blockCount = 1;
        copyInParams.blockLen = view.copyInSize * DATA_SIZE;
        copyInParams.srcStride = 0;
        copyInParams.dstStride = 0;
        AscendC::DataCopyPad(input, srcGlobal_[copyInOffset], copyInParams, padParams);
    }

    __aicore__ inline void CopyOutNonAligned(AscendC::LocalTensor<DataType> &output,
                                             uint32_t copyOutOffset, BroadcastDataView &view)
    {
        uint32_t blockId = AscendC::GetBlockIdx();
        uint32_t dstDataSize = view.outShape[0] * view.outShape[1];
        uint64_t dstShape = this->param_->tilingData.dstShape[1];
        AscendC::DataCopyExtParams copyOutParams;
        copyOutParams.blockLen = view.outShape[1] * DATA_SIZE;
        copyOutParams.blockCount = dstDataSize * DATA_SIZE / copyOutParams.blockLen;
        copyOutParams.srcStride = 0;
        if (view.outShape[1] + copyOutOffset % dstShape > dstShape) {
            // 列非对齐， 按实际数据拷贝
            copyOutParams.srcStride = OpsUtils::CeilAlign<uint32_t>(view.outShape[1], UB_ALIGN_COUNT) * DATA_SIZE;
            copyOutParams.blockLen = (dstShape - copyOutOffset % dstShape) * DATA_SIZE;
            copyOutParams.srcStride = (copyOutParams.srcStride - copyOutParams.blockLen) / ATVC::UB_ALIGN_32;
        }
        if (view.outShape[0] + copyOutOffset / dstShape > this->param_->tilingData.dstShape[0]) {
            // 行非对齐， 按实际数据拷贝
            copyOutParams.blockCount = (this->param_->tilingData.dstShape[0] - copyOutOffset / dstShape);
        }
        copyOutParams.dstStride = dstShape * DATA_SIZE - copyOutParams.blockLen;
        AscendC::DataCopyPad(dstGlobal_[copyOutOffset], output, copyOutParams);
    }

    __aicore__ inline void CopyOut(AscendC::LocalTensor<DataType> &output, uint32_t copyOutOffset, BroadcastDataView &view)
    {
        CopyOutNonAligned(output, copyOutOffset, view);
    }

    AscendC::GlobalTensor<DataType> srcGlobal_;
    AscendC::GlobalTensor<DataType> dstGlobal_;
    BroadcastCompute compute_;
    ParamPtr param_;
    KernelUtils::BroadcastBufPool bufPool_;
    constexpr static uint32_t DATA_SIZE = sizeof(DataType);
    constexpr static uint32_t UB_ALIGN_COUNT = ATVC::UB_ALIGN_32 / DATA_SIZE;
};
} // namespace Kernel
} // namespace ATVC
#endif // ATVC_BROADCAST_OP_TEMPLATE_H